# Agent for memory retrieval
# Written by Jiageng Mao
from agentdriver.llm_core.timeout import timeout
from agentdriver.memory.common_sense_memory import CommonSenseMemory
from agentdriver.memory.experience_memory import ExperienceMemory

class MemoryAgent:
    def __init__(self, data_path, model_name="gpt-3.5-turbo-0613", verbose=False, compare_perception=False) -> None:
        self.model_name = model_name
        self.common_sense_memory = CommonSenseMemory()
        self.experience_memory = ExperienceMemory(data_path, model_name=self.model_name, verbose=verbose, compare_perception=compare_perception)
        self.verbose = verbose

    def retrieve(self, working_memory):
        raise NotImplementedError
    
    def retrieve_common_sense_memory(self, knowledge_types: list = None):
        return self.common_sense_memory.retrieve(knowledge_types=knowledge_types)

    def retrieve_experience_memory(self, working_memory):
        return self.experience_memory.retrieve(working_memory)

    def insert(self, working_memory):
        raise NotImplementedError

    def update(self, working_memory):
        raise NotImplementedError

    @timeout(15)
    def run(self, working_memory):
        common_sense_prompts = self.retrieve_common_sense_memory()
        experience_prompt = self.retrieve_experience_memory(working_memory)
        return common_sense_prompts, experience_prompt